import pandas as pd
from sqlalchemy import create_engine
from urllib.parse import quote_plus
import aqicalc as aqi

stations = [
    'Aotizhongxin',
    'Changping',
    'Dingling',
    'Dongsi',
    'Guanyuan',
    'Gucheng',
    'Huairou',
    'Nongzhanguan',
    'Shunyi',
    'Tiantan',
    'Wanliu',
    'Wanshouxigong'
]

ranks = [51, 101, 151, 201, 301]


def to_iapi_limit(elem: int, cc: int, algo: str, limit: int) -> int:
    result = int(aqi.to_iaqi(elem, str(cc), algo=algo))
    if result > limit:
        return -1
    return result


def to_rank(aqi_value: int) -> int:
    for index, rank in enumerate(ranks):
        if aqi_value < rank:
            return index + 1
    return 6


engine = create_engine(f'mysql+pymysql://root:{quote_plus("123456")}@127.0.0.1:3306/air-quality')

# 定义你的SQL查询
sql_query = 'select * from t_air_quality_hour where station = "%s" order by `year`, `month`, `day`, `hour`'

for station in stations:
    df = pd.read_sql(sql_query % station, engine)
    df['pm25_iaqi'] = df['pm25'].apply(lambda x: int(aqi.to_iaqi(aqi.POLLUTANT_PM25, str(x), algo=aqi.ALGO_MEP)))
    df['pm25_24_avg_iaqi'] = df['pm25_24_avg'].apply(
        lambda x: int(aqi.to_iaqi(aqi.POLLUTANT_PM25, str(x), algo=aqi.ALGO_MEP)))
    df['pm10_iaqi'] = df['pm10'].apply(lambda x: int(aqi.to_iaqi(aqi.POLLUTANT_PM10, str(x), algo=aqi.ALGO_MEP)))
    df['pm10_24_avg_iaqi'] = df['pm10_24_avg'].apply(
        lambda x: int(aqi.to_iaqi(aqi.POLLUTANT_PM10, str(x), algo=aqi.ALGO_MEP)))
    df['so2_iaqi'] = df['so2'].apply(lambda x: to_iapi_limit(aqi.POLLUTANT_SO2_1H, x, aqi.ALGO_MEP, 800))
    df['no2_iaqi'] = df['no2'].apply(lambda x: int(aqi.to_iaqi(aqi.POLLUTANT_NO2_1H, str(x), algo=aqi.ALGO_MEP)))
    df['co_iaqi'] = df['co'].apply(lambda x: int(aqi.to_iaqi(aqi.POLLUTANT_CO_1H, str(x / 1000), algo=aqi.ALGO_MEP)))
    df['o3_iaqi'] = df['o3'].apply(lambda x: int(aqi.to_iaqi(aqi.POLLUTANT_O3_1H, str(x), algo=aqi.ALGO_MEP)))
    df['o3_8_avg_iaqi'] = df['o3_8_avg'].apply(lambda x: to_iapi_limit(aqi.POLLUTANT_O3_8H, x, aqi.ALGO_MEP, 800))

    df['aqi'] = df[
        ['pm25_iaqi', 'pm25_24_avg_iaqi', 'pm10_iaqi', 'pm10_24_avg_iaqi', 'so2_iaqi', 'no2_iaqi', 'co_iaqi', 'o3_iaqi',
         'o3_8_avg_iaqi']].max(axis=1)

    df['rank'] = df['aqi'].apply(to_rank)

    aqi_hour_df = df[
        ['year', 'month', 'day', 'hour', 'pm25_iaqi', 'pm25_24_avg_iaqi', 'pm10_iaqi', 'pm10_24_avg_iaqi', 'so2_iaqi',
         'no2_iaqi', 'co_iaqi', 'o3_iaqi', 'o3_8_avg_iaqi', 'aqi', 'rank', 'station']]

    aqi_hour_df.to_sql(name='t_aqi_hour', con=engine, if_exists='append', index=False)
